from docx import Document
from docx.shared import Pt, RGBColor
from docx.enum.text import WD_PARAGRAPH_ALIGNMENT
from docx.enum.table import WD_TABLE_ALIGNMENT
from utils.analysis.mediation.mediation_model4 import MediationModel4, Model4Result
import pandas as pd
from utils.report.base import format_table, create_table, delete_table_first_last_row, set_cell_border


def _generate_report(data, document, table_count):
    table1 = generate_table1(data, document, table_count)  # 表1
    table2 = generate_table2(data, document, table_count)  # 表2
    # 总结段落
    m2x_xz = "显著" if table1.cell(4, 1).text.endswith("*") else "不显著"
    m2x_way = "负向" if float(table1.cell(4, 2).text) < 0 else "正向"
    y2x_xz = "显著" if table1.cell(4, 3).text.endswith("*") else "不显著"
    y2x_way = "负向" if float(table1.cell(4, 4).text) < 0 else "正向"
    y2xm_xz1 = "显著" if table1.cell(4, 5).text.endswith("*") else "不显著"
    y2xm_way1 = "负向" if float(table1.cell(4, 6).text) < 0 else "正向"
    y2xm_xz2 = "显著" if table1.cell(5, 5).text.endswith("*") else "不显著"
    y2xm_way2 = "负向" if float(table1.cell(5, 6).text) < 0 else "正向"
    cover_str = ""
    if data.covar:
        cover_str += "，将"
        for i in range(len(data.covar)):
            cover_str += data.covar[i]
            if i != len(data.covar) - 1:
                cover_str += "、"
        cover_str += "作为控制变量"
    p2 = (
        f"本研究将{data.x}作为自变量{cover_str}，而将{data.m}作为中介变量进行分析，"
        f"从上表可以看出，首先，检验自变量{data.x}对中介变量{data.m}的直接影响，"
        f"由表可知{data.x}（β={table1.cell(4, 1).text},t={table1.cell(4, 2).text}）{m2x_xz}{m2x_way}影响中介变量{data.m}；"
        f"其次检验自变量{data.x}对因变量{data.y}的直接影响，"
        f"由表可知，自变量{data.x}（β={table1.cell(4, 3).text}，t={table1.cell(4, 4).text}）{y2x_xz}{y2x_way}影响因变量{data.y}；"
        f"最后，同时检验自变量{data.x}、中介变量{data.m}，对因变量{data.y}的影响，"
        f"由表可知，自变量{data.x}（β={table1.cell(4, 5).text}，t={table1.cell(4, 6).text}）{y2xm_xz1}{y2xm_way1}影响因变量{data.y}，"
        f"中介变量{data.m}（β={table1.cell(5, 5).text}，t={table1.cell(5, 6).text}）{y2xm_xz2}{y2xm_way2}影响因变量{data.y}。"
        f"由此可说明，{data.m}在{data.x}与{data.y}的关系间扮演部分中介作用。"
        f"即该总效应（{table2.cell(2, 1).text}），直接效应（{table2.cell(3, 1).text}）和间接效应（{table2.cell(4, 1).text}）。"
        f"针对{data.x}⇒{data.m}⇒{data.y}这条中介路径来看，中介效应值为{table2.cell(4, 1).text}，且95%区间并不包括数字0，因而说明此条中介效应路径存在。"
    )
    document.add_paragraph(p2)  # 插入段落
    return


def generate_table2(data, document, table_count):
    document.add_paragraph('')  # 插入段落
    # 新建表格
    init_rows = 5
    init_cols = 5
    # table = document.add_table(rows=init_rows, cols=init_cols)
    table = create_table(document, init_rows, init_cols)
    # 1. 填充1-2行
    # 填充第一行表头
    row = table.rows[0]  # 表格
    f_cell = None
    for i in range(init_cols):  # 合并单元格
        if i == 0:
            f_cell = table.cell(0, 0)
        else:
            f_cell.merge(table.cell(0, i))
    frequency_table_name = ""
    row.cells[0].text = frequency_table_name + '间接效应分析'
    table.cell(1, 0).text = "中介类型"
    table.cell(2, 0).text = "总效应"
    table.cell(3, 0).text = "直接效应"
    table.cell(4, 0).text = "间接效应"
    table.cell(1, 1).text = "Effect"
    table.cell(1, 2).text = "Boot SE"
    table.cell(1, 3).text = "BootLLCI"
    table.cell(1, 4).text = "BootULCI"
    tables = data.mediation_table
    for i in tables:
        if 'Total' == i.path:
            cur_row = 2
            table.cell(cur_row, 1).text = convert_str_to_float_3(i.coef)
            table.cell(cur_row, 2).text = convert_str_to_float_3(i.se)
            table.cell(cur_row, 3).text = convert_str_to_float_3(i.CI_2_5)
            table.cell(cur_row, 4).text = convert_str_to_float_3(i.CI_97_5)
        if 'Direct' == i.path:
            cur_row = 3
            table.cell(cur_row, 1).text = convert_str_to_float_3(i.coef)
            table.cell(cur_row, 2).text = convert_str_to_float_3(i.se)
            table.cell(cur_row, 3).text = convert_str_to_float_3(i.CI_2_5)
            table.cell(cur_row, 4).text = convert_str_to_float_3(i.CI_97_5)
        if 'Indirect' == i.path:
            cur_row = 4
            table.cell(cur_row, 1).text = convert_str_to_float_3(i.coef)
            table.cell(cur_row, 2).text = convert_str_to_float_3(i.se)
            table.cell(cur_row, 3).text = convert_str_to_float_3(i.CI_2_5)
            table.cell(cur_row, 4).text = convert_str_to_float_3(i.CI_97_5)
    # 格式化表格
    format_table(table)
    add_three_lines_table2(table)
    return table


def add_three_lines_table1(table):
    delete_table_first_last_row(table)
    first_row = table.rows[1]
    second_row = table.rows[3]
    for cell in first_row.cells:
        set_cell_border(cell, top={"sz": 12, "val": "single", "color": "#000000", "space": "0"})
    for cell in second_row.cells:
        set_cell_border(cell, top={"sz": 4, "color": "#000000", "val": "single", "space": "0"})
    last_row = table.rows[-1]
    for cell in last_row.cells:
        set_cell_border(cell, top={"sz": 12, "val": "single", "color": "#000000", "space": "0"})
    return


def add_three_lines_table2(table):
    delete_table_first_last_row(table)
    first_row = table.rows[1]
    for cell in first_row.cells:
        set_cell_border(cell, top={"sz": 12, "val": "single", "color": "#000000", "space": "0"},
                        bottom={"sz": 4, "color": "#000000", "val": "single", "space": "0"})
    last_row = table.rows[-1]
    for cell in last_row.cells:
        set_cell_border(cell, bottom={"sz": 12, "val": "single", "color": "#000000", "space": "0"})
    return


def generate_table1(data, document, table_count):
    # 新建表格
    init_rows = 6
    init_cols = 7
    # table = document.add_table(rows=init_rows, cols=init_cols, style="Light Shading")
    table = create_table(document, init_rows, init_cols)
    # 1. 填充1-3行
    # 填充第一行表头
    row = table.rows[0]  # 表格
    f_cell = None
    for i in range(init_cols):  # 合并单元格
        if i == 0:
            f_cell = table.cell(0, 0)
        else:
            f_cell.merge(table.cell(0, i))
    frequency_table_name = ""
    row.cells[0].text = frequency_table_name + '中介效应模型汇总'
    table.cell(1, 0).merge(table.cell(2, 0))
    table.cell(1, 0).text = "变量"
    table.cell(1, 1).merge(table.cell(1, 2))
    # m <- x
    table.cell(1, 1).text = f"模型1;{data.m}"
    table.cell(2, 1).text = "β"
    table.cell(2, 2).text = "t"
    table.cell(1, 3).merge(table.cell(1, 4))
    # y <- x
    table.cell(1, 3).text = f"模型2;{data.y}"
    table.cell(2, 3).text = "β"
    table.cell(2, 4).text = "t"
    table.cell(1, 5).merge(table.cell(1, 6))
    # y <- x + m
    table.cell(1, 5).text = f"模型3;{data.y}"
    table.cell(2, 5).text = "β"
    table.cell(2, 6).text = "t"
    # 2. 填充4-6行
    table.cell(3, 0).text = "常数"
    table.cell(4, 0).text = f"{data.x}"
    table.cell(5, 0).text = f"{data.m}"
    m2x_r2 = 0
    m2x_adj_r2 = 0
    for i in data.r_m2x.details:
        name = i.names
        if 'Intercept' == name:
            m2x_r2 = i.r2
            m2x_adj_r2 = i.adj_r2
            cur_str = convert_str_to_float_3(i.coef) + add_p_value(i.pval)
            table.cell(3, 1).text = cur_str
            table.cell(3, 2).text = convert_str_to_float_3(i.T)
        if data.x == name:
            cur_str = convert_str_to_float_3(i.coef) + add_p_value(i.pval)
            table.cell(4, 1).text = cur_str
            table.cell(4, 2).text = convert_str_to_float_3(i.T)
    y2x_r2 = 0
    y2x_adj_r2 = 0
    for i in data.r_y2x.details:
        name = i.names
        if 'Intercept' == name:
            y2x_r2 = i.r2
            y2x_adj_r2 = i.adj_r2
            cur_str = convert_str_to_float_3(i.coef) + add_p_value(i.pval)
            table.cell(3, 3).text = cur_str
            table.cell(3, 4).text = convert_str_to_float_3(i.T)
        if data.x == name:
            cur_str = convert_str_to_float_3(i.coef) + add_p_value(i.pval)
            table.cell(4, 3).text = cur_str
            table.cell(4, 4).text = convert_str_to_float_3(i.T)
    y2xm_r2 = 0
    y2xm_adj_r2 = 0
    for i in data.r_y2xm.details:
        name = i.names
        if 'Intercept' == name:
            y2xm_r2 = i.r2
            y2xm_adj_r2 = i.adj_r2
            cur_str = convert_str_to_float_3(i.coef) + add_p_value(i.pval)
            table.cell(3, 5).text = cur_str
            table.cell(3, 6).text = convert_str_to_float_3(i.T)
        if data.x == name:
            cur_str = convert_str_to_float_3(i.coef) + add_p_value(i.pval)
            table.cell(4, 5).text = cur_str
            table.cell(4, 6).text = convert_str_to_float_3(i.T)
        if data.m == name:
            cur_str = convert_str_to_float_3(i.coef) + add_p_value(i.pval)
            table.cell(5, 5).text = cur_str
            table.cell(5, 6).text = convert_str_to_float_3(i.T)
    # 3.填充控制变量
    covars = data.covar
    for i in covars:
        add_values(table, i, data.r_m2x.details, data.r_y2x.details, data.r_y2xm.details)
    # 4.填充R方，调整后的r方和F值
    add_last_3_cols(table, 'R2', convert_str_to_float_3(m2x_r2), convert_str_to_float_3(y2x_r2),
                    convert_str_to_float_3(y2xm_r2))
    add_last_3_cols(table, '调整R2', convert_str_to_float_3(m2x_adj_r2), convert_str_to_float_3(y2x_adj_r2),
                    convert_str_to_float_3(y2xm_adj_r2))
    add_last_3_cols(table, 'F值', convert_str_to_float_3(data.r_m2x.f) + add_p_value(data.r_m2x.p),
                    convert_str_to_float_3(data.r_y2x.f) + add_p_value(data.r_y2x.p),
                    convert_str_to_float_3(data.r_y2xm.f) + add_p_value(data.r_y2xm.p),
                    )
    # 格式化表格
    format_table(table)

    table.add_row()  # 表格动态增加一行
    cur_row = table.rows[-1]
    cur_row.cells[0].text = '* p<0.05 ** p<0.01 *** p<0.001'
    cur_row.cells[0].merge(cur_row.cells[1]).merge(cur_row.cells[2]).merge(cur_row.cells[3]).merge(
        cur_row.cells[4]).merge(cur_row.cells[5]).merge(cur_row.cells[6])
    add_three_lines_table1(table)
    return table


def add_values(table, cur_name, details_m2x, details_y2x, details_y2xm):
    table.add_row()  # 表格动态增加一行
    cur_row = table.rows[-1]
    cur_row.cells[0].text = cur_name
    for i in details_m2x:
        if cur_name == i.names:
            cur_str = convert_str_to_float_3(i.coef) + add_p_value(i.pval)
            cur_row.cells[1].text = cur_str
            cur_row.cells[2].text = convert_str_to_float_3(i.T)
    for i in details_y2x:
        if cur_name == i.names:
            cur_str = convert_str_to_float_3(i.coef) + add_p_value(i.pval)
            cur_row.cells[3].text = cur_str
            cur_row.cells[4].text = convert_str_to_float_3(i.T)
    for i in details_y2xm:
        if cur_name == i.names:
            cur_str = convert_str_to_float_3(i.coef) + add_p_value(i.pval)
            cur_row.cells[5].text = cur_str
            cur_row.cells[6].text = convert_str_to_float_3(i.T)


def add_last_3_cols(table, name, v1, v2, v3):
    table.add_row()  # 表格动态增加一行
    cur_row = table.rows[-1]
    cur_row.cells[0].text = name
    cur_row.cells[1].merge(cur_row.cells[2]).text = v1
    cur_row.cells[3].merge(cur_row.cells[4]).text = v2
    cur_row.cells[5].merge(cur_row.cells[6]).text = v3


def convert_str_to_float_3(value) -> str:
    if not value:
        return ''
    if isinstance(value, str):
        return "{:.3f}".format(float(value))
    elif isinstance(value, float):
        return "{:.3f}".format(value)
    return "{:.3f}".format(value)


def add_p_value(p):
    res = ''
    if p <= 0.001:
        res = '***'
    elif p <= 0.01:
        res = '**'
    elif p <= 0.05:
        res = '*'
    return res


def generate(src_data: [Model4Result], doc: Document = None) -> Document():
    if not doc:
        doc = Document()
    title1 = doc.add_heading(level=1)  # 增加标题
    t1_run = title1.add_run('中介效应分析-模型4')
    t1_run.font.color.rgb = RGBColor(10, 10, 10)
    for i in src_data:
        _generate_report(i, doc, 1)  # 生成表格
    return doc


if __name__ == '__main__':
    # 创建 Document 对象，等价于在电脑上打开一个 Word 文档
    doc = Document()
    pd.set_option('expand_frame_repr', False)
    # df = pd.read_excel('./1_24_test.xlsx')
    # # df = df[["性别", "年龄", "学历", "企业角色", 'JG', 'YY', 'JX', 'XW']]
    # # 1.此文件的结果不包含整体的F值和p值
    # # 2.若有多个x,y,m，则需要遍历进行获取结果
    # x_arr = ['JG']
    # y_arr = ['YY']
    # m_arr = ['XW']
    # covar = ["性别", "年龄", "学历", "企业角色"]

    df = pd.read_excel('./1-24-test.xlsx')
    x_arr = ['X1', 'X2', 'X3', 'X4', 'X5']
    y_arr = ['Y']
    m_arr = ['M']
    covar = []
    obj = MediationModel4(df, x_arr, y_arr, m_arr, covar)
    src_data = obj.analysis()
    title1 = doc.add_heading(level=1)  # 增加标题
    t1_run = title1.add_run('中介效应分析-模型4')
    p1 = ("如表所示，通过Process 宏程序中的 Bootstrap 方法, 结合逐步检验法，运用model 4进行中介作用分析，针对“M”这一中介变量，"
          "c表示X对Y时的回归系数（模型中没有中介变量M时），即总效应，当总效应显著时，讨论中介效应影响机制；a表示X对M时的回归系数，"
          "b表示M对Y时的回归系数，a*b为a与b的乘积即中介效应；c’表示X对Y时的回归系数（模型中有中介变量时），即直接效应；95% BootCI表示"
          "Bootstrap抽样计算得到的95%置信区间。如果a和b显著，c’不显著，且a*b的95% BootCI包括数字不包括0，则为完全中介；后续开始与"
          "文章背景开始结合；如果a和b显著，c’显著，且a*b的95% BootCI包括数字不包括0，则为部分中介；如果a和b至少一个不显著，"
          "且a*b的95% BootCI包括数字0（不显著），则中介作用不显著；")
    doc.add_paragraph(p1)  # 插入段落
    t1_run.font.color.rgb = RGBColor(10, 10, 10)
    for i in src_data:
        _generate_report(i, doc, 1)  # 生成表格
    # 保存文档
    doc.save('demo_zj_model4.docx')
